import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from torch import optim
from torch.utils.data import DataLoader, TensorDataset

# 解决中文的显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

"""
加载数据并制作数据集
"""


def make_Dataset(DataPath, batch_size):
    """加载数据并制作数据集"""
    # 读取原始数据
    df_data = pd.read_csv(DataPath, index_col=0)  # 将第一列设为索引
    df_data = df_data.fillna(method='pad')  # 缺失值填充：将空值用前一个值填充

    # 划分训练集特征和验证集特征
    data = np.array(df_data).flatten().tolist()  # 将数据转为列表
    data = pd.DataFrame(data)  # 再将数据转为DataFrame（本质上是对数据排列方式做了变换）
    train_data = data.iloc[:int(data.shape[0] * 0.8), [0]]  # 取前80%最为训练集
    val_data = data.iloc[int(data.shape[0] * 0.8):int(data.shape[0] * 0.9), [0]]  # 取80%~90%作为验证集

    # 特征数据归一化【训练集自己做自己的归一化，验证集也是自己做自己的归一化，并不会相互影响】
    scaler = MinMaxScaler(feature_range=(0, 1))  # 初始化MinMaxScaler，将数据归一化到[0, 1]
    train_data = scaler.fit_transform(train_data)
    val_data = scaler.transform(val_data)

    # 制作训练特征集和对应的训练标签集
    feature_num = 96
    x_train = []
    y_train = []
    for i in range(feature_num, len(train_data)):  # 从96个采样点开始，取前96个采样点作为特征，后一个采样点作为标签
        x_train.append(train_data[i - feature_num:i])
        y_train.append(train_data[i])
    x_train, y_train = np.array(x_train), np.array(y_train)  # 将训练集由list格式变为array格式

    # 制作验证特征集和对应的验证标签集
    x_val = []
    y_val = []
    for i in range(feature_num, len(val_data)):  # 从96个采样点开始，取前96个采样点作为特征，后一个采样点作为标签
        x_val.append(val_data[i - feature_num:i])
        y_val.append(val_data[i])
    x_val, y_val = np.array(x_val), np.array(y_val)  # 将训练集由list格式变为array格式

    # 将训练集和验证集转为tensor格式
    x_train = torch.tensor(x_train, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.float32)
    x_val = torch.tensor(x_val, dtype=torch.float32)
    y_val = torch.tensor(y_val, dtype=torch.float32)

    # 使用TensorDataset来封装数据集，便于后续使用Dataloader来加载数据
    train_dataset = TensorDataset(x_train, y_train)
    val_dataset = TensorDataset(x_val, y_val)

    # 使用Dataloader来批量加载数据
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)

    return train_loader, val_loader
